//===- Hello.cpp - Example code from "Writing an LLVM Pass" ---------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements two versions of the LLVM "Hello World" pass described
// in docs/WritingAnLLVMPass.html
//
//===----------------------------------------------------------------------===//

#include <llvm/Support/CommandLine.h>
#include <llvm/IRReader/IRReader.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/Support/SourceMgr.h>
#include <llvm/Support/ManagedStatic.h>
#include <llvm/IR/LegacyPassManager.h>
#include <llvm/Support/ToolOutputFile.h>

#include <llvm/Transforms/Scalar.h>
#include <llvm/Transforms/Utils.h>

#include <llvm/IR/Function.h>
#include <llvm/Pass.h>
#include <llvm/Support/raw_ostream.h>

#include <llvm/Bitcode/BitcodeReader.h>
#include <llvm/Bitcode/BitcodeWriter.h>

#include "llvm/IR/PassManager.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/IR/Module.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Utils/Mem2Reg.h"

#include <map>
#include <vector>
#include <string>
#include <set>
#include "DF4C.hh"

using namespace llvm;
using namespace std;
using namespace df4c;

namespace df4c {
  static bool EnableStderr = false;
  llvm::raw_ostream &errs() {
    return EnableStderr ? llvm::errs() : llvm::nulls();
  }
}

static ManagedStatic<LLVMContext> GlobalContext;
static LLVMContext &getGlobalContext() { return *GlobalContext; }

class EnableFunctionOptPass : public PassInfoMixin<EnableFunctionOptPass> {
public:
    PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM) {
        bool Changed = false;
        for (auto & F : M) {
          if (F.hasFnAttribute(Attribute::OptimizeNone)) {
              F.removeFnAttr(Attribute::OptimizeNone);
              Changed = true;
          }
        }
        return Changed ? PreservedAnalyses::none() 
                       : PreservedAnalyses::all();
    }
    
    static StringRef name() { return "EnableFunctionOptPass"; }
};

class FuncPtrPass : public PassInfoMixin<FuncPtrPass> {
private:
    // Map a function to its caller
    map<Function*, CallInst*> funccaller;
    // Map a line number to the set of called functions
    map<unsigned int, set<string>> linecallfunc;

    void visitCallInst(CallInst *callinst, vector<CallInst *> callinstList,
      int isAsValue = 0) {
      callinstList.push_back(callinst);
      Value * calledValue = callinst->getCalledOperand();
      visitValue(calledValue, callinstList, isAsValue);
    }

    void visitValue(Value *val, vector<CallInst *> callinstList,
      int isAsValue = 0) {
      if (Function *func = dyn_cast<Function>(val)) {
      // If a function is called
        if (isAsValue) {
          // If the value is a function (maybe "%1 = @fx(...)"? I guess)
          visitFunction(func, callinstList, isAsValue + 1);  // the "+1" is useless
        } else {
          visitFunction(func, callinstList);
        }
      } else if (PHINode *phi = dyn_cast<PHINode>(val)) {
        for (unsigned int i = 0u; i < phi->getNumIncomingValues(); ++i)
          visitValue(phi->getIncomingValue(i), callinstList, isAsValue);
      } else if (Argument *arg = dyn_cast<Argument>(val)) {
          visitArgument(arg, callinstList, isAsValue);
      } else if (CallInst *ci = dyn_cast<CallInst>(val)) {
          visitCallInst(ci, callinstList, isAsValue + 1);
      }
      else
        df4c::errs() << Normal_error << "Unknown value: " << *val << '\n';
    }

    void visitFunction(Function *func, vector<CallInst *> callinstList,
      int isAsValue = 0) {
      CallInst *ci = callinstList.back();
      callinstList.pop_back();

      // Map the function to the instruction
      funccaller[func] = ci;

      if (!isAsValue) {
        // Record a result!
        string fname = func->getName().str();
        unsigned int lineno = ci->getDebugLoc().getLine();
        linecallfunc[lineno].insert(fname);
      }

      for (BasicBlock &BB : *func)
        for (Instruction &inst : BB)
          if (!isa<DbgInfoIntrinsic>(inst)) {
            if (isAsValue) {  // Find return instructions
              if (ReturnInst * retinst = dyn_cast<ReturnInst>(&inst)) {
                Value * retVal = retinst->getReturnValue();
                visitValue(retVal, callinstList);
              }
            } else {  // Find call instructions
              if (CallInst * callinst = dyn_cast<CallInst>(&inst))
                visitCallInst(callinst, callinstList);
            }
          }

      funccaller.erase(func);
    }

    void visitArgument(Argument *arg, vector<CallInst *> callinstList, 
      int isAsValue = 0) {
      // If a value is an argument, first find the function that owns it
      Function * pfunc = arg->getParent();
      // and the index of the argument
      unsigned int argidx = arg->getArgNo();
      // Now, who called this function?
      if (funccaller.find(pfunc) != funccaller.end()) {
        CallInst * ci = funccaller[pfunc];
        // For each actual parameter passed in the call instruction
        Value * actarg = ci->getArgOperand(argidx);
        visitValue(actarg, callinstList, isAsValue);
      } else {
        df4c::errs() << Normal_error << "Cannot find caller for function: "
               << pfunc->getName() << '\n';
      }
    }

    void printResults() {
      for (const auto & l : linecallfunc) {
        outs() << l.first << " : ";
        bool flag_first = true;
        for (const auto & f : l.second) {
          if (!flag_first)
            outs() << ", ";
          else
            flag_first = false;
          outs() << f;
        }
        outs() << '\n';
      }
    }

public:
    static StringRef name() { return "FuncPtrPass"; }

    PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM) {
      /*
      outs() << "Hello: ";
      outs().write_escaped(M.getName()) << '\n';
      M.print(llvm::outs(), nullptr);
      outs()<<"------------------------------\n";
      */
      vector<CallInst *> callinstList;
      
      linecallfunc.clear();
      for (Function &F : M) {
        if (!F.isDeclaration()) {
          for (BasicBlock &BB : F) {
            for (Instruction &inst : BB) {
              if (!isa<DbgInfoIntrinsic>(inst)) {
                if (CallInst * callinst = dyn_cast<CallInst>(&inst)) {
                  visitCallInst(callinst, callinstList);
                }
              }
            }
          }
        }
      }
      printResults();
      return PreservedAnalyses::all();
    }


};

static cl::opt<std::string>
InputFilename(cl::Positional,
              cl::desc("<filename>.bc"),
              cl::init(""));
// Use "--enerr" or "-e" to enable stderr output
static cl::opt<bool> StdErrOption("enerr", llvm::cl::desc("Enable stderr output (for debug)"));
static cl::alias StdErrOptionShort("e", llvm::cl::aliasopt(StdErrOption));

int main(int argc, char **argv) {
  // This is for debug
  df4c::errs() << "Hello, LLVM!\n";

  LLVMContext &Context = getGlobalContext();
  SMDiagnostic Err;
  // Parse the command line to read the Inputfilename
  cl::ParseCommandLineOptions(argc, argv,
                              "FuncPtrPass \n My first LLVM too which does not do much.\n");
                            
  df4c::EnableStderr = StdErrOption;

  // Load the input module
  std::unique_ptr<Module> M = parseIRFile(InputFilename, Err, Context);
  if (!M) {
      Err.print(argv[0], df4c::errs());
      return 1;
  }

  llvm::PassBuilder PB;
  llvm::LoopAnalysisManager LAM;
  llvm::FunctionAnalysisManager FAM;
  llvm::CGSCCAnalysisManager CGAM;
  llvm::ModuleAnalysisManager MAM;

  PB.registerModuleAnalyses(MAM);
  PB.registerCGSCCAnalyses(CGAM);
  PB.registerFunctionAnalyses(FAM);
  PB.registerLoopAnalyses(LAM);
  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);

  llvm::ModulePassManager MPM;

  MPM.addPass(EnableFunctionOptPass());
  // The new mem2reg
  MPM.addPass(createModuleToFunctionPassAdaptor(llvm::PromotePass()));
  MPM.addPass(FuncPtrPass());

  MPM.run(*M, MAM);
  return 0;
}
